package cc.decat.escript.antlr4;

import cc.decat.escript.antlr4.gen.EScriptBaseVisitor;
import cc.decat.escript.antlr4.gen.EScriptLexer;
import cc.decat.escript.antlr4.gen.EScriptParser;
import cc.decat.escript.ir.AssignableNode;
import cc.decat.escript.ir.ChunkNode;
import cc.decat.escript.ir.Node;
import cc.decat.escript.ir.expression.*;
import cc.decat.escript.ir.statement.*;
import org.antlr.v4.runtime.tree.ParseTree;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

public class EScriptBaseVisitorImpl extends EScriptBaseVisitor<Node> {
    @Override
    public Node visitChunk(EScriptParser.ChunkContext ctx) {
        if(ctx.statementList() != null) {
            return new ChunkNode(visitStatementList(ctx.statementList()));
        }
        return new ChunkNode(null);
    }
    @Override
    public Node visitStatementList(EScriptParser.StatementListContext ctx) {
        return new StatementList(ctx.statement().stream().map(this::visit).collect(Collectors.toList()));
    }

    @Override
    public Node visitStatement(EScriptParser.StatementContext ctx) {
        return visit(ctx.getChild(0));
    }

    @Override
    public Node visitIdentifierAssignable(EScriptParser.IdentifierAssignableContext ctx) {
        AssignableNode assignableNode = new AssignableNode();
        assignableNode.setType(AssignableNode.Type.Variable);
        assignableNode.setVariableName(ctx.Identifier().getText());
        return assignableNode;
    }

    @Override
    public Node visitListIndexAssignable(EScriptParser.ListIndexAssignableContext ctx) {
        AssignableNode assignableNode = new AssignableNode();
        assignableNode.setType(AssignableNode.Type.ListIndex);
        assignableNode.setListExpression(visit(ctx.expression(0)));
        assignableNode.setIndexExpression(visit(ctx.expression(1)));
        return assignableNode;
    }

    @Override
    public Node visitMemberAttributeAssignable(EScriptParser.MemberAttributeAssignableContext ctx) {
        AssignableNode assignableNode = new AssignableNode();
        assignableNode.setType(AssignableNode.Type.MemberAttribute);
        assignableNode.setObjectExpression(visit(ctx.expression()));
        assignableNode.setAttributeName(ctx.Identifier().getText());
        return assignableNode;
    }

    @Override
    public Node visitAssignStatement(EScriptParser.AssignStatementContext ctx) {
        List<Node> assignableList = ctx.assignable().stream().map(this::visit).collect(Collectors.toList());
        List<Node> expressionList = ctx.expression().stream().map(this::visit).collect(Collectors.toList());
        return new AssignStatement(assignableList, expressionList);
    }

    @Override
    public Node visitReturnStatement(EScriptParser.ReturnStatementContext ctx) {
        return new ReturnStatement(ctx.expression() != null ? visit(ctx.expression()) : null);
    }

    @Override
    public Node visitIdentifierExpression(EScriptParser.IdentifierExpressionContext ctx) {
        return new VariableExpression(ctx.Identifier().getText());
    }

    @Override
    public Node visitLiteralExpression(EScriptParser.LiteralExpressionContext ctx) {
        return visitLiteral(ctx.literal());
    }

    @Override
    public Node visitLiteral(EScriptParser.LiteralContext ctx) {
        if(ctx.getChild(0) == ctx.StringLiteral()) {
            String s = ctx.StringLiteral().getText();
            return new LiteralExpression().setStringValue(s.substring(1, s.length() - 1));
        } else if(ctx.right == ctx.numberLiteral()) {
            LiteralExpression node = (LiteralExpression)visitNumberLiteral(ctx.numberLiteral());
            if(ctx.op == null || ctx.op.getType() != EScriptLexer.Minus) return node;
            if(node.getType() == LiteralExpression.Type.Integer) {
                return node.setIntegerValue(-node.getIntegerValue());
            } else if(node.getType() == LiteralExpression.Type.Float) {
                return node.setFloatValue(node.getFloatValue());
            }
        }
        return visit(ctx.getChild(0));
    }

    @Override
    public Node visitNumberLiteral(EScriptParser.NumberLiteralContext ctx) {
        if(ctx.getChild(0) == ctx.FloatLiteral()) {
            return new LiteralExpression().setFloatValue(Double.parseDouble(ctx.FloatLiteral().getText()));
        }
        return this.visitIntegerLiteral(ctx.integerLiteral());
    }

    @Override
    public Node visitIntegerLiteral(EScriptParser.IntegerLiteralContext ctx) {
        LiteralExpression le = new LiteralExpression();
        if(ctx.getChild(0) == ctx.HexIntegerLiteral()) {
            le.setIntegerValue(Long.parseLong(ctx.HexIntegerLiteral().getText(), 16));
        } else if(ctx.getChild(0) == ctx.DecimalLiteral()) {
            le.setIntegerValue(Long.parseLong(ctx.DecimalLiteral().getText(), 10));
        } else if (ctx.getChild(0) == ctx.OctalIntegerLiteral()) {
            le.setIntegerValue(Long.parseLong(ctx.OctalIntegerLiteral().getText(), 8));
        } else if (ctx.getChild(0) == ctx.BinaryIntegerLiteral()) {
            le.setIntegerValue(Long.parseLong(ctx.BinaryIntegerLiteral().getText(), 2));
        } else {
            le = null;
        }
        return le;
    }

    @Override
    public Node visitAdditiveExpression(EScriptParser.AdditiveExpressionContext ctx) {
        Node left = visit(ctx.expression(0));
        Node right = visit(ctx.expression(1));
        BinaryExpression.OpType opType = null;
        if(ctx.op.getType() == EScriptLexer.Plus) {
            opType = BinaryExpression.OpType.ADD;
        } else if(ctx.op.getType() == EScriptLexer.Minus) {
            opType = BinaryExpression.OpType.SUB;
        }
        return new BinaryExpression(left, opType, right);
    }

    @Override
    public Node visitMultiplicativeExpression(EScriptParser.MultiplicativeExpressionContext ctx) {
        Node left = visit(ctx.expression(0));
        Node right = visit(ctx.expression(1));
        BinaryExpression.OpType opType = null;
        if(ctx.op.getType() == EScriptLexer.Multiply) {
            opType = BinaryExpression.OpType.MUL;
        } else if(ctx.op.getType() == EScriptLexer.Divide) {
            opType = BinaryExpression.OpType.DIV;
        } else if(ctx.op.getType() == EScriptLexer.Modulus) {
            opType = BinaryExpression.OpType.MOD;
        }
        return new BinaryExpression(left, opType, right);
    }

    @Override
    public Node visitRelationalExpression(EScriptParser.RelationalExpressionContext ctx) {
        Node left = visit(ctx.expression(0));
        Node right = visit(ctx.expression(1));
        BinaryExpression.OpType opType = null;
        if(ctx.op.getType() == EScriptLexer.LessThan) {
            opType = BinaryExpression.OpType.LT;
        } else if(ctx.op.getType() == EScriptLexer.LessThanEquals) {
            opType = BinaryExpression.OpType.LTE;
        } else if(ctx.op.getType() == EScriptLexer.GreaterThan) {
            opType = BinaryExpression.OpType.GT;
        } else if(ctx.op.getType() == EScriptLexer.GreaterThanEquals) {
            opType = BinaryExpression.OpType.GTE;
        }
        return new BinaryExpression(left, opType, right);
    }

    @Override
    public Node visitParenthesizedExpression(EScriptParser.ParenthesizedExpressionContext ctx) {
        return visit(ctx.expression());
    }

    @Override
    public Node visitBlockStatement(EScriptParser.BlockStatementContext ctx) {
        return new BlockStatement(ctx.statementList() != null ? visitStatementList(ctx.statementList()) : new StatementList());
    }

    @Override
    public Node visitCallFunctionStatement(EScriptParser.CallFunctionStatementContext ctx) {
        return new CallFunctionStatement(visit(ctx.expression()), ctx.expressionList() != null ? ctx.expressionList().expression().stream().map(this::visit).collect(Collectors.toList()) : Collections.emptyList());
    }

    @Override
    public Node visitCallFunctionExpression(EScriptParser.CallFunctionExpressionContext ctx) {
        return new CallFunctionExpression(visit(ctx.expression()), ctx.expressionList() != null ? ctx.expressionList().expression().stream().map(this::visit).collect(Collectors.toList()) : Collections.emptyList());
    }

    @Override
    public Node visitFunctionDeclarationStatement(EScriptParser.FunctionDeclarationStatementContext ctx) {
        FunctionDeclarationStatement functionDeclarationStatement = new FunctionDeclarationStatement();
        functionDeclarationStatement.setName(ctx.Identifier().getText());
        functionDeclarationStatement.setParams(ctx.parameters() != null ? ctx.parameters().Identifier().stream().map(ParseTree::getText).collect(Collectors.toList()) : Collections.emptyList());
        functionDeclarationStatement.setBlock(visitBlockStatement(ctx.blockStatement()));
        return functionDeclarationStatement;
    }

    @Override
    public Node visitAnonymousFunctionExpression(EScriptParser.AnonymousFunctionExpressionContext ctx) {
        return new AnonymousFunctionExpression(ctx.parameters() != null ? ctx.parameters().Identifier().stream().map(ParseTree::getText).collect(Collectors.toList()) : Collections.emptyList(), visitBlockStatement(ctx.blockStatement()));
    }

    @Override
    public Node visitIfStatement(EScriptParser.IfStatementContext ctx) {
        IfStatement ifStatement = new IfStatement();
        ifStatement.setCondExpr(new ArrayList<>()); ifStatement.getCondExpr().add(visit(ctx.expression(0)));
        ifStatement.setThenBody(new ArrayList<>()); ifStatement.getThenBody().add(visitBlockStatement(ctx.blockStatement(0)));
        for(int i = 1; i < ctx.expression().size(); i++) {
            ifStatement.getCondExpr().add(visit(ctx.expression(i)));
            ifStatement.getThenBody().add(visitBlockStatement(ctx.blockStatement(i)));
        }
        ifStatement.setElseBody(ctx.else_ != null ? visitBlockStatement(ctx.else_) : null);
        return ifStatement;
    }

    @Override
    public Node visitForStatement(EScriptParser.ForStatementContext ctx) {
        ForStatement forStatement = new ForStatement();
        forStatement.setInitialization(ctx.init != null ? visitAssignStatement(ctx.init) : null);
        forStatement.setCondition(ctx.expression() != null ? visit(ctx.expression()) : null);
        forStatement.setIncrement(ctx.inc != null ? visitAssignStatement(ctx.inc) : null);
        forStatement.setBody(visitBlockStatement(ctx.blockStatement()));
        return forStatement;
    }

    @Override
    public Node visitLogicAndExpression(EScriptParser.LogicAndExpressionContext ctx) {
        return new BinaryExpression(visit(ctx.expression(0)), BinaryExpression.OpType.LOGIC_AND, visit(ctx.expression(1)));
    }

    @Override
    public Node visitLogicOrExpression(EScriptParser.LogicOrExpressionContext ctx) {
        return new BinaryExpression(visit(ctx.expression(0)), BinaryExpression.OpType.LOGIC_OR, visit(ctx.expression(1)));
    }

    @Override
    public Node visitListExpression(EScriptParser.ListExpressionContext ctx) {
        return new ListExpression(ctx.expressionList() != null ? ctx.expressionList().expression().stream().map(this::visit).collect(Collectors.toList()) : Collections.emptyList());
    }

    @Override
    public Node visitListIndexExpression(EScriptParser.ListIndexExpressionContext ctx) {
        return new ListIndexExpression(visit(ctx.expression(0)), visit(ctx.expression(1)));
    }

    @Override
    public Node visitDictExpression(EScriptParser.DictExpressionContext ctx) {
        DictExpression dictExpression = new DictExpression();
        if(ctx.dictItems() == null) {
            return dictExpression;
        }
        for(int i = 0; i < ctx.dictItems().expression().size(); i += 2) {
            dictExpression.getMap().put(visit(ctx.dictItems().expression(i)), visit(ctx.dictItems().expression(i + 1)));
        }
        return dictExpression;
    }

    @Override
    public Node visitEqualityExpression(EScriptParser.EqualityExpressionContext ctx) {
        Node left = visit(ctx.expression(0));
        Node right = visit(ctx.expression(1));
        BinaryExpression.OpType opType = null;
        if(ctx.op.getType() == EScriptLexer.Equals) {
            opType = BinaryExpression.OpType.EQ;
        } else if(ctx.op.getType() == EScriptLexer.NotEquals) {
            opType = BinaryExpression.OpType.NEQ;
        }
        return new BinaryExpression(left, opType, right);
    }

    @Override
    public Node visitContinueStatement(EScriptParser.ContinueStatementContext ctx) {
        return new ContinueStatement();
    }

    @Override
    public Node visitBreakStatement(EScriptParser.BreakStatementContext ctx) {
        return new BreakStatement();
    }

    @Override
    public Node visitMemberExpression(EScriptParser.MemberExpressionContext ctx) {
        return new MemberExpression(visit(ctx.expression()), ctx.Identifier().getText());
    }
}
